# -*- coding: utf-8 -*-
# @Time     : 2020/8/20 3:32 下午
# @Author   : huangxiong
# @FileName : holding_analysis.py
# @Comment  : 组合持仓分析相关的基础函数
# @Software : PyCharm
import pandas as pd

from quant_researcher.quant.datasource_fetch.fund_api.fund_holding_related import get_holding_report_info
from quant_researcher.quant.performance_attribution.core_functions.holding_analysis.allocation_analysis.holding_based_allocation import \
    get_asset_allocation
from quant_researcher.quant.datasource_fetch.index_api.index_components import get_index_component_weights
from quant_researcher.quant.project_tool import time_tool
from quant_researcher.quant.datasource_fetch.portfolio_api import portfolio_tool
from quant_researcher.quant.project_tool.logger.my_logger import LOG
from quant_researcher.quant.datasource_fetch.index_api.index_constant import SW_CODE_DICT, SW_INDUSTRY_CODE_TRANS_DICT, CSRC_CODE_DICT
from quant_researcher.quant.datasource_fetch.stock_api.stock_info import get_stock_industry_classification


def get_portfolio_asset_allocation(portfolio_id, calc_date=None):
    """
    计算组合某一天的大类资产配置情况，默认最新一天

    :param portfolio_id: 组合ID
    :param calc_date: 计算日期，格式"20200810"
    :return: pd.DataFrame
    """

    if calc_date is None:
        calc_date = time_tool.get_yesterday(marker='without_n_dash')

    fund_weight_df = portfolio_tool.get_portfolio_fund_weight(portfolio_id, calc_date)
    if fund_weight_df is None:
        LOG.error(f"{portfolio_id} 该组合id有问题")
        return
    fund_list = fund_weight_df['fund_code'].tolist()

    f_asset_df = get_asset_allocation(fund_list, net_or_total='t', only_latest=False,
                                      end_date=calc_date)
    f_asset_df = f_asset_df.merge(fund_weight_df, how='left', on='fund_code')
    f_asset_df = f_asset_df.dropna()
    p_asset_df = f_asset_df.iloc[:, 1:5].mul(f_asset_df['fund_weight'], axis=0).sum().reset_index()
    p_asset_df = p_asset_df.rename(columns={'index': 'asset_type', 0: 'weight'})
    return p_asset_df


def get_portfolio_industry_allocation(portfolio_id, calc_date=None, industry_type='sw2014'):
    """
    计算组合的行业占比

    :param portfolio_id: 组合ID
    :param calc_date: 计算日期，格式"20200810"
    :param industry_type: 行业类型，sw2014-申万，csrc2012-证监会，默认申万
    :return: pd.DataFrame
    """
    if calc_date is None:
        calc_date = time_tool.get_yesterday(marker='without_n_dash')

    fund_weight_df = portfolio_tool.get_portfolio_fund_weight(portfolio_id, calc_date)
    if fund_weight_df is None:
        LOG.error(f"{portfolio_id} 该组合id有问题")
        return
    fund_list = fund_weight_df['fund_code'].tolist()

    f_industry_df = get_latest_industry_allocation(fund_list, calc_date, industry_type=industry_type)

    if industry_type == 'sw2014':
        indu_code_dict = SW_CODE_DICT
    else:
        indu_code_dict = CSRC_CODE_DICT

    if f_industry_df is None:
        p_industry_df = pd.Series(indu_code_dict).reset_index()
        p_industry_df.columns = ['industry_code', 'industry_name']
        p_industry_df['weight'] = 0
    else:
        f_industry_df = f_industry_df.merge(fund_weight_df, how='left', on='fund_code')
        f_industry_df = f_industry_df.dropna()
        f_industry_df['fund_weight'] = f_industry_df['fund_weight'] / f_industry_df['fund_weight'].sum()
        p_industry_df = f_industry_df[list(indu_code_dict.keys())].mul(f_industry_df['fund_weight'], axis=0).sum().reset_index()
        p_industry_df = p_industry_df.rename(columns={'index': 'industry_code', 0: 'weight'})
        p_industry_df['industry_name'] = p_industry_df['industry_code'].replace(indu_code_dict)
    return p_industry_df


def get_latest_industry_allocation(asset_list, end_date, asset_type='fund',
                                   detail_stock=False, industry_type='sw2014'):
    """
    得到一个基金列表或指数列表的每个基金或指数在计算日期前最新的行业配置数据

    :param asset_list: 计算的基金或指数列表
    :param end_date: 计算日期，格式"20200821"
    :param asset_type: 计算的基金或指数类型，可选fund、index，默认fund
    :param detail_stock: 是否返回每支股票的占比和行业分类，默认False
    :param industry_type: 行业类型，sw2014-申万，csrc2012-证监会，默认申万
    :return:
    """
    if asset_type == 'fund':
        stock_df = get_holding_report_info(report='stock', fund_code=asset_list, report_type=[5, 6],
                                           end_date=end_date, select=['fund_code', 'stock_code', 'hold_mktvalue',
                                                                      'pptinnv as stock_weight'])
        stock_df['stock_code'] = stock_df['stock_code'].str.extract(r'(\d+)')
        stock_df['stock_code'] = stock_df['stock_code'].str.zfill(5)
        stock_df = stock_df.dropna()
        stock_df['stock_weight'] = stock_df['stock_weight'] / 100
    else:
        stock_df = get_index_component_weights(asset_list, end_date=end_date)
        if stock_df is None:
            return
        stock_df = stock_df.rename(columns={'end_date': 'trade_date'})
        stock_df['trade_date'] = stock_df['trade_date'].apply(lambda x: x.replace('-', ''))

    if stock_df.empty:
        LOG.error(f"未找到{asset_list}在{end_date}之前的股票持仓信息")
        return stock_df

    # industry_stock = get_index_hist_stock_pool(SW_CODE_LIST, end_date, end_date)
    industry_stock = get_stock_industry_classification(induclsstd_code=industry_type, end_date=end_date)
    if industry_stock is None:
        return
    industry_stock = industry_stock.rename(columns={'indu_code': 'industry_code'})
    industry_stock = industry_stock.drop_duplicates(subset=['stock_code'])
    if industry_type == 'sw2014':
        industry_stock['industry_code'] = industry_stock['industry_code'].replace(SW_INDUSTRY_CODE_TRANS_DICT)
        indu_code_list = list(SW_CODE_DICT.keys())
    else:
        indu_code_list = list(CSRC_CODE_DICT.keys())

    industry_df = stock_df.merge(industry_stock, how='left', on='stock_code')
    if detail_stock:
        return industry_df
    else:
        if asset_type == 'fund':
            industry_df = industry_df.groupby(['fund_code', 'industry_code'])['stock_weight'].sum().unstack()
        else:
            industry_df = industry_df.groupby(['index_code', 'industry_code'])['stock_weight'].sum().unstack()
        industry_df = industry_df.reindex(columns=indu_code_list)
        industry_df = industry_df.fillna(0)
        return industry_df


if __name__ == '__main__':
    # 获取基金行业配置权重测试
    test_res = get_latest_industry_allocation(asset_list=['110022'], end_date='2020-01-01', asset_type='fund')
    # 获取基金资产配置权重测试
    # test_res = get_fund_latest_asset_allocation(fund_list=['110022'], calc_date='20200821',
    #                                             n_or_t='t')

